package zjut.raymond.witdemo.algorithm;

import android.graphics.PointF;

public class UnscentKalmanFilter {
    private PointF pdrLocation;
    private PointF wifiLocation;
    private final static int STATE_DIMENSION = 2; //状态量维数
    private final static int MEASUREMENT_DIMENSION = 2; //观测量维数
    private final static double alpha = 0.001d;
    private final static double bate = 2d;
    //  private static double lamda = alpha * alpha * STATE_DIMENSION - STATE_DIMENSION;
    private static double lamda = 3 - STATE_DIMENSION;
    private static double weight_0_m = lamda / (STATE_DIMENSION + lamda);
    // private static double weight_0_c = lamda / (STATE_DIMENSION + lamda) + (1 - alpha * alpha + bate);
    private static double weight_0_c = lamda / (STATE_DIMENSION + lamda);
    private static double weight_i = 1 / (2 * (STATE_DIMENSION + lamda));

    private static double e_pdr = 0.1d * 36;        //PDR平均误差
    private static double e_wifi = 0.76d * 36;       //Wifi定位平均误差
    private static double[][] Q = {{1.0 * 36, 0}, {0, 1.0 * 36}};   //预测噪声
    private static double[][] Q_x = {{e_pdr, 0}, {0, e_pdr}};   //预测噪声
//    private static double[][] V_k = {{e_wifi}, {e_wifi}};   //观测噪声
    private static double[] V_k = {e_wifi, e_wifi};   //观测噪声
    private static double[][] noise_R = {{e_wifi, 0}, {0, e_wifi}};

    private static double[][] X_average;
    private static double[][] Z_average;
    private static double[][] K;

    public interface MyCallBack {
        void onLocationChange(PointF location);
    }

    private MyCallBack callBack;

    public void setCallBack(MyCallBack callBack) {
        this.callBack = callBack;
    }

    public UnscentKalmanFilter(PointF pdrLocation) {
        this.pdrLocation = pdrLocation;
    }

    public PointF getPdrLocation() {
        return pdrLocation;
    }

    public void onStepUpdate(double stepLength, double orientation) {
        //计算协方差矩阵的cholesky分解:
        double[][] L = JamaUtils.MatrixCholesky(Q);
        double[][] x_init = {{pdrLocation.x}, {pdrLocation.y}};
        //生成sigma点
        double[][] sigma = CreateSigmaPoint(x_init, L, STATE_DIMENSION, lamda);

        //将选出的sigma点带入非线性函数f得到Sigma点集chi:chi_i = f(sigma_0_i)：
        double[][] X_i = new double[STATE_DIMENSION][2 * STATE_DIMENSION + 1];
        for (int i = 0; i < 2 * STATE_DIMENSION + 1; i++) {
            double[][] chi_i = Non_Linear_f(JamaUtils.getColumnByIndex(sigma, i), stepLength, orientation);
            X_i = JamaUtils.InsertMatrixByIndex(X_i, chi_i, i);
        }

        //利用chi集求解预测状态和及其方差：
        X_average = CalculateAverage(X_i, weight_0_m, weight_i);

        double[][] P_X_variance = CalculateVariance(X_i, X_average, weight_0_c, weight_i, Q_x);

        double[][] cholesky_P_variance = JamaUtils.MatrixCholesky(P_X_variance);
        double[][] z_sigma = CreateSigmaPoint(X_average, cholesky_P_variance, STATE_DIMENSION, lamda);
        //将新的sigma点带入非线性函数h得到z_i:z_i = h(z_sigma_i)：
        double[][] Z_i = Non_Linear_h(z_sigma, MEASUREMENT_DIMENSION, V_k);
        //  ---------------------计算观测量均值-------------------------;

        Z_average = CalculateAverage(Z_i, weight_0_m, weight_i);
        //  JamaUtils.dispMatrix(Z_average);

        //  System.out.println("---------------------计算观测量方差-------------------------");

        double[][] P_Z_variance = CalculateVariance(Z_i, Z_average, weight_0_c, weight_i, noise_R);

        //计算状态和测量之间的协方差：
        //   System.out.println("---------------------计算协方差-------------------------");

        double[][] P_X_Z_variance = CalculateCovariance(X_i, X_average, Z_i, Z_average, weight_0_c, weight_i);
        //  JamaUtils.dispMatrix(P_X_Z_variance);

        //计算卡尔曼增益：
        //  System.out.println("---------------------计算观测量逆矩阵-------------------------");

        double[][] P_Z_variance_inverse = JamaUtils.matrixInverse(P_Z_variance);
        K = JamaUtils.matrixMult(P_X_Z_variance, P_Z_variance_inverse);

        pdrLocation.x = (float) X_average[0][0];
        pdrLocation.y = (float) X_average[1][0];
        callBack.onLocationChange(pdrLocation);
    }

    public void onWifiLocationChange(PointF wifiLocation) {
        if (K != null && K.length > 0) {
            double[][] z_measure = {{wifiLocation.x}, {wifiLocation.y}};
            double[][] X_res = UpdateState(X_average, Z_average, z_measure, K);
            pdrLocation.x = (float) X_res[0][0];
            pdrLocation.y = (float) X_res[1][0];
            callBack.onLocationChange(pdrLocation);
        }
    }

    public static double[][] CreateSigmaPoint(double[][] x_0, double[][] L, int n, double lamda) {
        double[][] res = new double[x_0.length][2 * n + 1];
        res = JamaUtils.InsertMatrixByIndex(res, x_0, 0);
        double[][] temp;
        for (int i = 1; i < 2 * n + 1; i++) {
            temp = JamaUtils.getColumnByIndex(L, (i - 1) % n);
            temp = JamaUtils.matrixMultiply(temp, Math.sqrt(n + lamda));
            if (i <= n) {
                temp = JamaUtils.matrixAdd(x_0, temp);
            } else {
                temp = JamaUtils.matrixSub(x_0, temp);
            }
            res = JamaUtils.InsertMatrixByIndex(res, temp, i);
        }
        return res;
    }

    private static double[][] Non_Linear_f(double[][] m, double step_length, double orientation) {
        double[][] res = new double[JamaUtils.getMatrixRow(m)][JamaUtils.getMatrixColumn(m)];
        res[0][0] = m[0][0] + step_length * Math.sin(Math.toRadians(orientation));
        res[1][0] = m[1][0] - step_length * Math.cos(Math.toRadians(orientation));
        return res;
    }

    //    public static double[][] Non_Linear_h(double[][] z_sigma, int n_z, double x_0, double y_0, double x_wifi, double y_wifi, double[][] v_k) {
//        double[][] res = new double[n_z][JamaUtils.getMatrixColumn(z_sigma)];
//        for (int i = 0; i < JamaUtils.getMatrixColumn(z_sigma); i++) {
//            res[0][i] = z_sigma[0][i] + v_k[0][0];
//            res[1][i] = z_sigma[1][i] + v_k[1][0];
//            res[2][i] = Math.sqrt(Math.pow(z_sigma[0][i] - x_0, 2) + Math.pow(z_sigma[1][i] - y_0, 2)) + v_k[2][0];
//            res[3][i] = z_sigma[2][i] + v_k[3][0];
//        }
//        return res;
//    }
    public static double[][] Non_Linear_h(double[][] z_sigma, int n_z, double[] v_k) {
        double[][] res = new double[n_z][JamaUtils.getMatrixColumn(z_sigma)];
        for (int i = 0; i < JamaUtils.getMatrixColumn(z_sigma); i++) {
            res[0][i] = z_sigma[0][i] + v_k[0];
            res[1][i] = z_sigma[1][i] + v_k[1];
        }
        return res;
    }

    public static double[][] CalculateAverage(double[][] X, double omega_0, double omega_i) {
        double[][] res = new double[JamaUtils.getMatrixRow(X)][1];
        double[][] temp;
        for (int i = 0; i < JamaUtils.getMatrixColumn(X); i++) {
            temp = JamaUtils.getColumnByIndex(X, i);
            if (i == 0) {
                temp = JamaUtils.matrixMultiply(temp, omega_0);
            } else {
                temp = JamaUtils.matrixMultiply(temp, omega_i);
            }
            res = JamaUtils.matrixAdd(res, temp);
        }
        return res;
    }

    public static double[][] CalculateVariance(double[][] X, double[][] x_bar, double omega_0, double omega_i, double[][] noise) {
        double[][] res = new double[x_bar.length][x_bar.length];
        double[][] var_temp_1, var_temp_2;
        for (int i = 0; i < JamaUtils.getMatrixColumn(X); i++) {
            var_temp_1 = JamaUtils.matrixSub(JamaUtils.getColumnByIndex(X, i), x_bar);

            var_temp_2 = JamaUtils.matrixMult(var_temp_1, JamaUtils.matrixTranspose(var_temp_1));

            if (i == 0) {
                var_temp_2 = JamaUtils.matrixMultiply(var_temp_2, omega_0);
            } else {
                var_temp_2 = JamaUtils.matrixMultiply(var_temp_2, omega_i);
            }

            res = JamaUtils.matrixAdd(res, var_temp_2);
        }
        res = JamaUtils.matrixAdd(res, noise);
        return res;
    }

    public static double[][] CalculateCovariance(double[][] X, double[][] x_bar, double[][] Z, double[][] z_bar, double omega_0, double omega_i) {
        double[][] res = new double[x_bar.length][z_bar.length];
        double[][] x_temp, z_temp, temp;
        for (int i = 0; i < JamaUtils.getMatrixColumn(X); i++) {
            x_temp = JamaUtils.matrixSub(JamaUtils.getColumnByIndex(X, i), x_bar);
            z_temp = JamaUtils.matrixSub(JamaUtils.getColumnByIndex(Z, i), z_bar);
            z_temp = JamaUtils.matrixTranspose(z_temp);
            temp = JamaUtils.matrixMult(x_temp, z_temp);
            if (i == 0) {
                temp = JamaUtils.matrixMultiply(temp, omega_0);
            } else {
                temp = JamaUtils.matrixMultiply(temp, omega_i);
            }
            res = JamaUtils.matrixAdd(res, temp);
        }
        return res;
    }


    public static double[][] UpdateState(double[][] x_bar, double[][] z_bar, double[][] z_measure, double[][] K) {
        double[][] res = JamaUtils.matrixSub(z_measure, z_bar);
        res = JamaUtils.matrixMult(K, res);
        res = JamaUtils.matrixAdd(x_bar, res);
        return res;
    }
}
